from motor.motor_asyncio import AsyncIOMotorClient
from beanie import init_beanie
import logging
from app.config import settings
from app.models.visa import (
    VisaApplication,
    DiscoursePost,
    ScrapedWebsite,
    VisaStatistics
)

logger = logging.getLogger(__name__)

class Database:
    client: AsyncIOMotorClient = None
    database = None

database = Database()

async def connect_to_mongodb():
    """连接到 MongoDB 数据库"""
    try:
        mongodb_url = settings.get_mongodb_url()
        logger.info(f"Connecting to MongoDB at {mongodb_url}")
        
        # 创建 MongoDB 客户端
        database.client = AsyncIOMotorClient(
            mongodb_url,
            maxPoolSize=10,
            minPoolSize=10,
            serverSelectionTimeoutMS=5000,
        )
        
        # 获取数据库
        database.database = database.client[settings.MONGODB_DATABASE]
        
        # 初始化 Beanie
        await init_beanie(
            database=database.database,
            document_models=[
                VisaApplication,
                DiscoursePost,
                ScrapedWebsite,
                VisaStatistics,
            ]
        )
        
        logger.info("Successfully connected to MongoDB and initialized Beanie")
        
    except Exception as e:
        logger.error(f"Error connecting to MongoDB: {e}")
        raise

async def close_mongodb_connection():
    """关闭 MongoDB 连接"""
    try:
        if database.client:
            database.client.close()
            logger.info("MongoDB connection closed")
    except Exception as e:
        logger.error(f"Error closing MongoDB connection: {e}")

async def ping_database():
    """检查数据库连接状态"""
    try:
        await database.client.admin.command('ping')
        return True
    except Exception as e:
        logger.error(f"Database ping failed: {e}")
        return False

# 数据库工具函数
async def create_indexes():
    """创建数据库索引"""
    try:
        # VisaApplication 索引
        await VisaApplication.get_motor_collection().create_index([("passport_number", 1)], unique=True)
        await VisaApplication.get_motor_collection().create_index([("status", 1)])
        await VisaApplication.get_motor_collection().create_index([("destination_country", 1)])
        await VisaApplication.get_motor_collection().create_index([("application_date", -1)])
        
        # DiscoursePost 索引
        await DiscoursePost.get_motor_collection().create_index([("post_id", 1)], unique=True)
        await DiscoursePost.get_motor_collection().create_index([("topic_id", 1)])
        await DiscoursePost.get_motor_collection().create_index([("created_at", -1)])
        await DiscoursePost.get_motor_collection().create_index([("related_countries", 1)])
        
        # ScrapedWebsite 索引
        await ScrapedWebsite.get_motor_collection().create_index([("url", 1)], unique=True)
        await ScrapedWebsite.get_motor_collection().create_index([("scraped_at", -1)])
        await ScrapedWebsite.get_motor_collection().create_index([("country", 1)])
        await ScrapedWebsite.get_motor_collection().create_index([("site_type", 1)])
        
        # VisaStatistics 索引
        await VisaStatistics.get_motor_collection().create_index([("country", 1), ("visa_type", 1)])
        await VisaStatistics.get_motor_collection().create_index([("period_start", 1), ("period_end", 1)])
        await VisaStatistics.get_motor_collection().create_index([("created_at", -1)])
        
        logger.info("Database indexes created successfully")
        
    except Exception as e:
        logger.error(f"Error creating indexes: {e}")

async def get_database_stats():
    """获取数据库统计信息"""
    try:
        stats = {
            "visa_applications": await VisaApplication.count(),
            "discourse_posts": await DiscoursePost.count(),
            "scraped_websites": await ScrapedWebsite.count(),
            "visa_statistics": await VisaStatistics.count(),
        }
        return stats
    except Exception as e:
        logger.error(f"Error getting database stats: {e}")
        return {}
